from utils.asynccmd import AsyncCmdMixin
import pipes
import asyncio
import json
from export.export import export_to_excel

filter_url = lambda x: x['url']
class HTTPXScan(AsyncCmdMixin):
    def __init__(self, path, domains: list):
        self._exec_path = path
        self.domains = domains
        # 存放存活的域名
        self.results = []

    @property
    def columns_dicts(self):
        return {
            'scheme': '协议',
            'port': '端口',
            'path': '请求路径',
            'body-sha256': '内容(sha256)',
            'header-sha256': '请求头(sha256)',
            'a': 'A记录',
            'title': '标题',
            'webserver': 'webserver',
            'host': '主机地址',
            'response-time': '响应时间'
        }

    async def scan(self, arguments=""):
        """TODO: 修改为批量执行"""
        results = []
        not_alive_results = []
        for x in self.domains:
            try:
                cmd = f'echo {x} | {self._exec_path} -silent -json'
                result = await self.run_cmd(cmd, is_json=True)
                if result:
                    try:
                        print(result)
                        result['a'] = ','.join(result['a'])
                        result.pop('timestamp')
                        results.append(result)
                    except Exception as e:
                        print(e)
                        result = None
                        print(f'- {x} 不存活')
                        not_alive_results.append(x)
            except Exception as e:
                print(f'- {x} 不存活')
                not_alive_results.append(x)
                continue
        self.results = results

    def export(self,is_save, save_path=None):
        if not is_save:
            return
        export_to_excel(self.results, sheet_name='存活域名', xlsx_name='test',
                        columns_dict=self.columns_dicts, save_path=save_path)


class HTTPX:
    def __init__(self, exec_path, domains):
        self.httpx = HTTPXScan(exec_path, domains)
        self.results = []

    def run(self, save_path, is_save=True):
        asyncio.run(self.httpx.scan())
        self.results = list(map(filter_url, self.httpx.results))
        self.httpx.export(is_save, save_path=save_path)
